import torch
from torch import nn
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

df_train_data = pd.read_csv('iris_train.data')
df_test_data = pd.read_csv('iris_test.data')

x = df_train_data.iloc[:, 0:4]
x = np.array(x, dtype=float)
x = torch.unsqueeze(torch.FloatTensor(x), dim=1)
print(x)

df_train_data['种类编码'] = df_train_data['种类'].map({'Iris-setosa': 1, 'Iris-versicolor': 2, 'Iris-virginica': 3})
df_test_data['种类编码'] = df_test_data['种类'].map({'Iris-setosa': 1, 'Iris-versicolor': 2, 'Iris-virginica': 3})

y = df_train_data['种类编码']

y = np.array(y, dtype=int)
y = torch.unsqueeze(torch.FloatTensor(y), dim=1)
print(y)
